// kernel/vm.c
#include "riscv.h"
#include "types.h"
#include "defs.h"
#include "param.h"
#include "spinlock.h"
#include "vm.h"
#include "pmm.h"
#include "memlayout.h"

static void map_region(pagetable_t pt, uint64 va, uint64 pa, uint64 size, int perm);
// 创建空页表
pagetable_t create_pagetable(void)
{
    pagetable_t pt = (pagetable_t)alloc_page();
    if(pt) {
        memset(pt, 0, PAGE_SIZE);
    }
    return pt;
}

// 释放页表（递归释放）
void free_pagetable(pagetable_t pt, int level)
{
    if((uint64)pt % PAGE_SIZE != 0 || pt == 0)
        panic("free_pagetable: invalid page table");
    
    // 递归释放下级页表
    if(level > 0) {
        for(int i = 0; i < 512; i++) {
            pte_t pte = pt[i];
            if(pte & PTE_V && (pte & (PTE_R|PTE_W|PTE_X)) == 0) {
                uint64 child = PTE2PA(pte);
                free_pagetable((pagetable_t)child, level - 1);
            }
        }
    }
    
    free_page(pt);
}

// 页表遍历（查找现有PTE）
pte_t* walk_lookup(pagetable_t pt, uint64 va)
{
    return walk_create(pt, va, 0);
}

// 页表遍历（必要时创建）
pte_t* walk_create(pagetable_t pt, uint64 va, int alloc)
{
    if(va >= MAXVA) panic("walk_create: virtual address too high");
    
    pagetable_t cur = pt;
    pte_t *pte;
    
    for(int level = 2; level >= 0; level--) {
        pte = &cur[PX(level, va)];
        if(*pte & PTE_V) {
            // 有效PTE，进入下一级
            cur = (pagetable_t)PTE2PA(*pte);
        } else {
            // PTE无效
            if(!alloc) return 0;
            
            // 分配新页表页
            pagetable_t new_pt = (pagetable_t)alloc_page();
            if(new_pt == 0) return 0;
            
            memset(new_pt, 0, PAGE_SIZE);
            *pte = PA2PTE((uint64)new_pt) | PTE_V;
            cur = new_pt;
        }
    }
    return pte;
}

// 建立页映射
int map_page(pagetable_t pt, uint64 va, uint64 pa, int perm)
{
    if(va % PAGE_SIZE != 0 || pa % PAGE_SIZE != 0)
        panic("map_page: unaligned address");
    
    pte_t *pte = walk_create(pt, va, 1);
    if(pte == 0) return -1;
    
    if(*pte & PTE_V) {
        panic("map_page: remapping existing page");
    }
    
    *pte = PA2PTE(pa) | perm | PTE_V;
    return 0;
}

// 解除页映射
int unmap_page(pagetable_t pt, uint64 va)
{
    pte_t *pte = walk_lookup(pt, va);
    if(pte == 0 || !(*pte & PTE_V)) {
        return -1; // 未映射
    }
    
    *pte = 0; // 清除PTE
    
    // 刷新TLB
    asm volatile("sfence.vma zero, %0" : : "r" (va));
    
    return 0;
}

// 调试：打印页表
void dump_pagetable(pagetable_t pt, int level)
{
    static const char* level_names[] = {"L2", "L1", "L0"};
    
    for(int i = 0; i < 512; i++) {
        pte_t pte = pt[i];
        if(pte & PTE_V) {
            // 打印缩进
            for(int l = 0; l < 2 - level; l++) cprintf("  ");
            
            cprintf("%s[%d]: PTE %p -> PA %p perm: ", 
                   level_names[level], i, pte, PTE2PA(pte));
            
            if(pte & PTE_R) cprintf("R");
            if(pte & PTE_W) cprintf("W"); 
            if(pte & PTE_X) cprintf("X");
            if(pte & PTE_U) cprintf("U");
            if(pte & PTE_G) cprintf("G");
            cprintf("\n");
            
            // 递归打印下级页表
            if(level > 0 && (pte & (PTE_R|PTE_W|PTE_X)) == 0) {
                dump_pagetable((pagetable_t)PTE2PA(pte), level - 1);
            }
        }
    }
}



pagetable_t kernel_pagetable;

// 创建内核页表
void kvminit(void)
{
    kernel_pagetable = create_pagetable();
    
    // 恒等映射内核代码和数据
    // 内核文本段：R+X权限
    map_region(kernel_pagetable, KERNBASE, KERNBASE, 
               (uint64)etext - KERNBASE, PTE_R | PTE_X);
    
    // 内核数据段：R+W权限  
    map_region(kernel_pagetable, (uint64)etext, (uint64)etext,
               PHYSTOP - (uint64)etext, PTE_R | PTE_W);
    
    // 映射设备
    map_region(kernel_pagetable, UART0, UART0, PAGE_SIZE, PTE_R | PTE_W);
    map_region(kernel_pagetable, VIRTIO0, VIRTIO0, PAGE_SIZE, PTE_R | PTE_W);
    
    cprintf("kvminit: kernel page table created\n");
}

// 映射连续内存区域
void map_region(pagetable_t pt, uint64 va, uint64 pa, uint64 size, int perm)
{
    if(size % PAGE_SIZE != 0)
        panic("map_region: size not page aligned");
    
    for(uint64 i = 0; i < size; i += PAGE_SIZE) {
        if(map_page(pt, va + i, pa + i, perm) < 0) {
            panic("map_region: failed to map page");
        }
    }
}

// SATP寄存器构造宏
//#define MAKE_SATP(pagetable) (8L << 60 | ((uint64)pagetable) >> 12)

// 激活内核页表
void kvminithart(void)
{
    w_satp(MAKE_SATP(kernel_pagetable));
    sfence_vma(); // 刷新TLB
    
    cprintf("kvminithart: virtual memory enabled\n");
}

// TLB刷新
// void sfence_vma()
// {
//     asm volatile("sfence.vma zero, zero");
// }
